import numpy as np
from PIL import Image
import io
try:
    from .typedef import point_t, spmTarget_t, spmOp_t, dpmOp_t, scan_t, rawpoint_t
    from .utilsCommon import pySva
except:
    from typedef import point_t, spmTarget_t, spmOp_t, dpmOp_t, scan_t, rawpoint_t
    from utilsCommon import pySva

ADC_PWR = 13 - 20 * np.log10(2**16) # 1v vpp 50ohm= 13dBm

def deCompress(M_in: bytearray, data_packing='CP4D', endians='big'):
    """de compress nxp`s cp8 or cp4 data

    Args:
        M_in (bytearray): input bytes of compress data
        data_packing (str, optional): CP4D or CP8. Defaults to 'CP4D'.
        endians (str, optional): endians. Defaults to 'big'.

    Returns:
        np.ndarray: decompress data
    """
    M_in = np.frombuffer(M_in, dtype='>Q' if endians == 'big' else 'Q')
    if len(M_in) % 8 != 0:
        print(f"Input data length{len(M_in)}illegality!")
        return None
    common = M_in >> 60 if endians == 'big' else (M_in & 0xF)
    common = common.astype('H')
    if data_packing == 'CP8':
        indices = np.arange(16)
        effBits = np.r_[3 * np.ones(2), 4 * np.ones(12), 3 * np.ones(2)].astype('H')
        mask = 2**effBits - 1
        mantissaPtr = 60 if endians == 'big' else 4
    else:
        indices = np.arange(8)
        effBits = 7 * np.ones(8, dtype='H')
        mask = 2**effBits - 1
        mantissaPtr = 56 if endians == 'big' else 8
    o = np.empty([len(effBits), len(M_in)], dtype='I')
    for e, m, i in zip(effBits, mask, indices):
        if endians == 'big':
            mantissaPtr -= e
            o[i, :] = ((M_in >> mantissaPtr) & m) << (32 - e)
        else:
            o[i, :] = ((M_in >> mantissaPtr) & m) << (32 - e)
            mantissaPtr += e
            
    o.dtype = np.int32
    o = o >> (8 + common)
    o = o.T.reshape(-1).astype(np.float32)
    o.dtype = np.complex64
    return o

def unpackScanInfo(infoBytes):
    lam = 3e8 / 77.2e9
    info = {}
    ptr = 0
    if len(infoBytes) != scan_t.itemsize:
        return info
    temp = np.frombuffer(infoBytes, dtype=scan_t)
    NeedSwap = temp[0]['ntx'] > 100  #ntx will not large than 100
    if NeedSwap: temp = temp.byteswap()
    ptr += scan_t.itemsize
    for i in scan_t.names: info[i] = temp[i][0]
    # =============  calc radar parameter ============ #
    info["deltaR"] = 3e8 / 2 / info["effbwKhz"] / 1e3
    info["offsetR"] = info["rstart"] * info["deltaR"]
    vmax = lam / 2 / info["tchirpns"] / 1e-9
    info["deltaV"] = vmax / info["chirps"]
    if info['nrx'] == 16:
        info["dBase"] = 0.56 if info['ntx'] == 12 else 0.9
    else:
        info['dBase'] = 0.6
    info['offsetRbin'] = info["rstart"]
    info['offsetVbin'] = info["vbins"] // 2
    info['rAxis'] = (np.arange(info['rbins']) + info['rstart']) * info['deltaR']
    info['vAxis'] = (np.arange(info['vbins']) - info["vbins"] // 2) * info["deltaV"]
    info['tAxis'] = (np.arange(info['abins']) - info['abins'] // 2) / info['abins'] / info['dBase']
    return info

def unpackCloudRdc3(infoBytes, RDC3Bytes, tsBase = 0):
    ret = []
    if (len(infoBytes) != 40):
        infoBytes = b"\x01\x00\x000\x01\x00\x10\x00\x00\x10\x00\x00\x00\x02\x00\x00\xc0'\t\x00Pp\x02\x00\x00\x05\x00\x00\x00\x02\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00"
    h_t = np.dtype({'names': ('iSize', 'iCapacity'), 'formats':('i', 'i')}, align=True)
    h = np.copy(np.frombuffer(RDC3Bytes[:h_t.itemsize], dtype=h_t))[0]
    points = np.frombuffer(RDC3Bytes[h_t.itemsize:], dtype=point_t)
    tsMsList = np.unique(points['iStampMs'])
    for tsMs in tsMsList:
        p = points[points['iStampMs'] == tsMs]
        h['iSize'] = len(p)
        ret.append({
            'stamp': tsBase + float(tsMs) / 1000,
            'point': p,
            'bytes': b'\x01\x00' + infoBytes[2:] + h.tobytes() + p.tobytes()
        })
    return ret

def unpackRdc3(InfoRdc3Bytes, infoInclude=True):
    h_t = np.dtype({'names': ('iSize', 'iCapacity'), 'formats':('i', 'i')}, align=True)
    if infoInclude:
        info = unpackScanInfo(InfoRdc3Bytes[:scan_t.itemsize])
        RDC3Bytes = InfoRdc3Bytes[scan_t.itemsize:]
        h = np.frombuffer(RDC3Bytes[:h_t.itemsize], dtype=h_t).copy()[0]
        points = np.frombuffer(RDC3Bytes[h_t.itemsize:], dtype=rawpoint_t)
        if h['iSize'] != len(points): points = points[:h['iSize']]
        return info, points
    else:
        h = np.frombuffer(InfoRdc3Bytes[:h_t.itemsize], dtype=h_t).copy()[0]
        points = np.frombuffer(InfoRdc3Bytes[h_t.itemsize:], dtype=rawpoint_t)
        if h['iSize'] != len(points): points = points[:h['iSize']]
        return None, points

def unpackPoint(point):
    header_size = 16
    if point is not np.ndarray: point = np.frombuffer(point, dtype='b')
    header = point[:header_size]
    detect = point[header_size:]
    header.dtype = 'i'
    iSize = header[3]
    if len(detect) != iSize * spmTarget_t.itemsize: return np.zeros(0, spmTarget_t)
    detect.dtype = spmTarget_t
    return detect

def unpackSpm(data):
    point = unpackPoint(data)
    d = np.empty(len(point), dtype=spmOp_t)
    d["velocity"] = point['fSpeed']
    d["range"] = point['fRange']
    d["azimuth"] = point["fAzangle"]
    d["elevation"] = point["fElangle"]
    d["SNR"] = point["fSNR"]
    d["power"] = point["fSNR"]
    d['fX'] = d['range'] * np.cos(d["elevation"]) * np.sin(d["azimuth"])
    d['fY'] = d['range'] * np.cos(d["elevation"]) * np.cos(d["azimuth"])
    d['fZ'] = d['range'] * np.sin(d["elevation"])
    d['rbin'] = point['peakIdx'] >> 16
    d['vbin'] = point['peakIdx'] & 0xffff
    d['peakIdx'] = -1
    return {"detection": d}

def unpackDpm(data):
    if data is not np.ndarray:
        data = np.frombuffer(data, dtype='b')
    if len(data) == 0:
        return {"track": None, "trace": None}
    header = data[:16]
    header.dtype = 'Q'
    iSize = int(header[1])
    obj = data[16:16 + iSize * dpmOp_t.itemsize]
    obj.dtype = dpmOp_t
    trace = data[16 + iSize * dpmOp_t.itemsize : ]
    trace.dtype = 'i'
    traceSize = trace[0]
    trace.dtype = 'f'
    trace = trace[1:]
    eol = np.argwhere(trace == -1000).flatten()
    traceList = []
    if len(trace) == traceSize and len(eol) == len(obj):
        eol = np.r_[-1, eol]
        for i in range(1, len(eol)):
            item = trace[eol[i-1]+1:eol[i]].reshape([-1, 2])
            traceList.append(item)
    return {"track": obj, "trace": traceList}

def parseImgToRdc3(imgBytes):
    img = np.array(Image.open(io.BytesIO(imgBytes)))
    idx = np.argwhere(img.any(axis=-1))
    sel = img[idx[:,0], idx[:,1]]
    rdc3 = np.zeros(len(idx), dtype=point_t)
    rdc3['iRangeBin'] = idx[:, 0]
    rdc3['iDopplerBin'] = idx[:, 1]
    rdc3['iSnrLog2'] = (sel[:,0].astype('H'))<<4
    rdc3['iBeamsBin'] = sel[:,1]
    rdc3['iDropFlag'] = sel[:,2]
    rdc3['iStampMs'] = sel[:,3]
    return rdc3

def unpackRdm(rdmBytes, info):
    rdm = np.frombuffer(rdmBytes, 'h')
    sft = rdm[:info['rbins']]
    rdm = rdm[info['rbins']:].reshape([info['rbins'],info['vbins']])
    rdm = (rdm.T-(512*sft)).T   # 归一化adptv shift
    rdm = np.fft.fftshift(rdm, axes=-1)
    rdm = rdm.astype('f') / 85
    rdm = ADC_PWR + rdm - info['gain'] # 归一化adptv shift
    return rdm

def unpackAdc(adcBytes, info):
    if info['nrx'] == 8:
        TILE = 8
        adc = np.frombuffer(adcBytes, 'h')
        adc = adc.byteswap()
        adc = adc.reshape([2, 2 ,info['samples'] // TILE, 4, TILE]).transpose([0, 3, 1, 2, 4]).reshape([8, info['samples'] * 2])
    else:
        TILE = 16
        adc = np.frombuffer(adcBytes, 'h')
        adc = adc.reshape([4, 2 ,info['samples'] // TILE, 4, TILE]).transpose([0, 3, 1, 2, 4]).reshape([16, info['samples'] * 2])
    return adc

def unpackRdc1(rdc1Bytes, info, NeedSwap=False):
    if len(rdc1Bytes) >= info['nrx'] * info['chirps'] * info['rbins'] * 4:
        CP = 'CMPLX16'
        cubesize = info['nrx'] * info['chirps'] * info['rbins'] * 4
    elif len(rdc1Bytes) >= info['nrx'] * info['chirps'] * info['rbins'] * 2:
        CP = 'CP4D'
        cubesize = info['nrx'] * info['chirps'] * info['rbins'] * 2
    else:
        CP = 'CP8'
        cubesize = info['nrx'] * info['chirps'] * info['rbins']
    if CP == 'CMPLX16':
        _fft1d = np.frombuffer(rdc1Bytes[:cubesize], dtype = 'h')
        if NeedSwap: _fft1d = _fft1d.byteswap()
        fft1d = np.zeros([info['rbins'], info['chirps'], info['nrx']], dtype=np.complex64)
        fft1d.real = _fft1d[1::2].reshape([info['rbins'], info['chirps'], info['nrx']])
        fft1d.imag = _fft1d[0::2].reshape([info['rbins'], info['chirps'], info['nrx']])
    else:
        ENDIAN = 'big' if NeedSwap else 'little'
        fft1d = rdc1Bytes[:cubesize]
        fft1d = deCompress(fft1d, CP, ENDIAN).reshape([info['rbins'], info['chirps'], info['nrx']])
    rSft = rdc1Bytes[cubesize:]
    rSft = np.frombuffer(rSft, dtype = 'h')
    if NeedSwap: rSft = rSft.byteswap()
    if np.max(rSft)>30:
        rSft = rSft.reshape([info['nrx'], info['chirps']])/32767
        rSft[6:,:] = rSft[3]
        rSft[4:6,:] = rSft[2]
        rSft[2:4,:] = rSft[1]
        rSft[0:2,:] = rSft[0]
    else:
        rSft = rSft.reshape([info['nrx']//2, info['chirps']]).T
        temp = rSft[:,0:0]
        for i in range(rSft.shape[1]):
            temp = np.c_[temp, rSft[:,i]]
            temp = np.c_[temp, rSft[:,i]]
        rSft = 1/2**(temp)
    fft1d = fft1d*(rSft)
    return fft1d

rdc_ext_t = np.dtype({'names': ('range', 'velocity','azimuth','elevation', 'snr', 'x', 'y', 'z'),'formats': ('f', 'f', 'f', 'f', 'f', 'f', 'f', 'f')}, align=True)
basic_fields = [
            ('range', {'units': 'm'}),
            ('velocity', {'units': 'm'}),
            ('azimuth', {'units': '°'}),
            ('elevation', {'units': '°'}),
            ('snr', {'units': 'dB'}),
            ('x', {'units': 'm'}),
            ('y', {'units': 'm'}),
            ('z', {'units': 'm'})]

def externRdc3(rdc3: np.ndarray, scan: dict) -> np.ndarray:
    """
    在 rdc3 基础上增加 rdc_ext_t 的 dtype, 并填充扩展字段的值。
    """
    num_points = rdc3.shape[0]
    extended_rdc3 = np.zeros(num_points, dtype=rdc3.dtype.descr + rdc_ext_t.descr)
    for field in rdc3.dtype.names: extended_rdc3[field] = rdc3[field]
    if 'iRangeBin' in rdc3.dtype.names and 'deltaR' in scan:
        extended_rdc3['range'] = rdc3['iRangeBin'] * scan['deltaR'] + scan["offsetR"]
    if 'iDopplerBin' in rdc3.dtype.names and 'deltaV' in scan:
        half = scan['vbins'] // 2
        val = rdc3['iDopplerBin']
        doppler_bin = np.where(val > half, val - scan['vbins'], val)
        extended_rdc3['velocity'] = doppler_bin * scan['deltaV']
    if 'iSnrLog2' in rdc3.dtype.names:
        extended_rdc3['snr'] = rdc3['iPwrLog2'] / 85
    elif 'iPwrLog2' in rdc3.dtype.names:
        extended_rdc3['snr'] = rdc3['iSnrLog2'] / 85
    if 'iAngleBin' in rdc3.dtype.names and 'range' in extended_rdc3.dtype.names:
        sin_azi = rdc3['iAngleBin'] / scan['abins'] / scan['dBase']
        azi = np.arcsin(sin_azi)
        extended_rdc3['azimuth'] = np.degrees(azi)
        extended_rdc3['x'] = extended_rdc3['range'] * np.cos(azi)
        extended_rdc3['y'] = extended_rdc3['range'] * sin_azi
        extended_rdc3['z'] = 0  # 假设 z 为 0
    return extended_rdc3
import os
__dir__ = os.path.dirname(os.path.realpath(__file__))
default_cal = np.load(__dir__ + '/cal_R202_new.npy')
def transOrg2image(org, cal=default_cal):
    nfft = 64
    try:
        org = deCompress(org, endians='big')
    except:
        org = None
    if org is None: return None
    org = np.dot(org.reshape([-1,48]), cal)
    org = np.fliplr(org)
    sarImag = np.fft.fft(org*np.sqrt(np.hamming(48)), nfft, axis=-1).astype(np.complex64)
    sarImag = pySva(sarImag, axis=0)
    sarImag = np.fft.fftshift(sarImag, axes=1) / nfft
    ram = 20 * np.log10(np.abs(sarImag))
    return sarImag, ram
